##Function to calculate the number and percent of thalamocortical connections showing a significant association with an environmental characteristic
sig.effects <- function(df, env.measure){
  enveffects.df <- df
  enveffects.df$significant <- p.adjust(enveffects.df$GAM.cov.pvalue, method = c("fdr")) < 0.05
  sigeffect.totaln <- sum(enveffects.df$significant)
  sigeffect.percent <- round(sigeffect.totaln/length(enveffects.df$significant)*100, 2)
  print(sprintf("%s thalamocortical connections (%s percent) show a significant association between thalamocortical FA and the %s", sigeffect.totaln, sigeffect.percent, env.measure))
}

spin.data <- left_join(spin.parcels, sigeffects.df, by = "tract") #merged data for spin tests, with all 360 parcels in the expected right hemisphere --> left hemisphere order (matching glasser.coords)
spin.data <- merge(spin.data, S.A.axis, by = c("label", "orig_parcelname", "tract"), sort =F)

cor.test(spin.data$GAM.cov.tvalue, spin.data$SA.axis, method = c("spearman"))

perm.sphere.p(x = spin.data$SA.axis, y = spin.data$GAM.cov.tvalue, perm.id = perm.id.full, corr.type = "spearman") 

ggplot() +
  geom_point(data = spin.data, aes(x = SA.axis, y = GAM.cov.tvalue, color = SA.axis), size = 0.5) +
  geom_smooth(data = spin.data, aes(x = SA.axis, y = GAM.cov.tvalue), method = "lm", color = c("black"), formula = y ~ x, linewidth = 0.2) + 
  scale_color_gradient2(low = "goldenrod1", mid = "#ede4f5", high = "#6f1282", guide = "colorbar", aesthetics = "color", midpoint = 180) +
  theme_classic() +
  labs(x="\nS-A axis", y="Environment effect (t)") +
  theme(
  legend.position = "none", 
  axis.text = element_text(size = 5, family = "Arial", color = c("black")),
  axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
  axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
  axis.line = element_line(linewidth = .2), 
  axis.ticks = element_line(linewidth = .2)) 

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/EnvEffect_SArank_correlationplot.pdf", device = "pdf", dpi = 500, width = 1.6, height = 1.3)

#Glasser parcel x,y,z coordinates
load(file = "/cbica/projects/thalamocortical_development/Maps/parcellations/surface/hcp_mmp1.0.rda") #RDA file from the brainGraph package https://github.com/cwatson/brainGraph/blob/master/data/hcp_mmp1.0.rda with glasser region x, y, and z MNI coordinates. regions are in lh --> rh order
hcp_mmp1.0$x.mni[1:180]  <- hcp_mmp1.0$x.mni[1:180]*-1 #convert R-> L hemisphere coords into medial lateral coords
hcp_mmp1.0$z.mni <- hcp_mmp1.0$z.mni*-1 
hcp_mmp1.0$label <- NA
hcp_mmp1.0$label[1:180] <- glasser.regions$label[181:360]
hcp_mmp1.0$label[181:360] <- glasser.regions$label[1:180]
hcp_mmp1.0 <- hcp_mmp1.0 %>% select(-hemi)

#Function to compute the correlation between a environment map and an anatomical axis (x,y,z) as well as test the significance of the difference between the correlation with the anatomical axis versus the S-A axis
compute_axis_correlation <- function(axis){
  df <- merge(spin.data, hcp_mmp1.0, by = "label", sort = F)
  completeobs.n <- sum(!is.na(df$GAM.cov.tvalue))
  
  #S-A axis - environment correlation (as computed above) for comparison of two overlapping correlations based on dependent groups
  SA.axis.cor <-  cor(df$SA.axis, df$GAM.cov.tvalue, use = "complete.obs", method = c("spearman"))
  
  #Anatomical axis - environment correlation
  anatomical.axis.cor <- cor(df[,axis], df$GAM.cov.tvalue, use = "complete.obs", method = c("spearman")) #correlation between anatomical axis and env metric
  anatomical.axis.pvalue <- perm.sphere.p(x = df[,axis], y = df$GAM.cov.tvalue, perm.id = perm.id.full, corr.type = "spearman") #spin based p-value for the correlation
  
  #S-A axis - anatomical axis correlation
  SA.anatomical.cor <- cor(df$SA.axis, df[,axis], use = "complete.obs", method = c("spearman"))
  
  #Test for significant difference between two correlations based on dependent groups (i.e., correlations with an overlapping input)
  cocor.result <- cocor.pvalue <- cocor.dep.groups.overlap(
  r.jk = SA.axis.cor, 
  r.jh = anatomical.axis.cor, 
  r.kh = SA.anatomical.cor, 
  n = completeobs.n, alternative = "two.sided", test = "hittner2003")
  
  cocor.pvalue <- cocor.result@hittner2003[["p.value"]]
  
  comparison.results <- data.frame(axis = axis, axis.cor = anatomical.axis.cor, axis.cor.pvalue = anatomical.axis.pvalue, SA.cor = SA.axis.cor, cocor.pvalue = cocor.pvalue)
  
  return(comparison.results)
  }

ggseg(.data = hcp_mmp1.0, atlas = "glasser", mapping = aes(fill = y.mni)) + paletteer::scale_fill_paletteer_c("harrypotter::lunalovegood")

compute_axis_correlation(axis = "y.mni")

ggseg(.data = hcp_mmp1.0, atlas = "glasser", mapping = aes(fill = x.mni)) + paletteer::scale_fill_paletteer_c("harrypotter::lunalovegood")

compute_axis_correlation(axis = "x.mni")

ggseg(.data = hcp_mmp1.0, atlas = "glasser", mapping = aes(fill = z.mni)) + paletteer::scale_fill_paletteer_c("harrypotter::lunalovegood")

compute_axis_correlation(axis = "z.mni")

anatomical.ps <- c(0.0005890574, 0.3319559, 0.3909002)
p.adjust(anatomical.ps, method = c("fdr"))

#Put above results into df for plotting
axis.results <- data.frame(Axis = factor(), corr = double())
axis.results <- axis.results %>% add_row(Axis = "S-A", corr = 0.3010667)
axis.results <- axis.results %>% add_row(Axis = "A-P", corr = 0.06846274)
axis.results <- axis.results %>% add_row(Axis = "D-V", corr = 0.2117744)
axis.results <- axis.results %>% add_row(Axis = "M-L", corr = 0.1841123)
axis.results$Axis <- factor(axis.results$Axis, ordered = TRUE, levels = c("S-A", "D-V", "M-L", "A-P"))

ggplot(axis.results, aes(x = Axis, y = corr, fill = Axis)) + 
  geom_bar(stat='identity') + 
  labs(x="") +
  labs(y="\nEnvironment Effect Correlation") +
  theme_classic()+
  scale_fill_manual(values=c(alpha("#323280", 0.5), "#672975", "#672975", "#672975")) +
  theme(
  legend.position = "none", 
  axis.text = element_text(size = 5, family = "Arial", color = c("black")),
  axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
  axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
  axis.line = element_line(linewidth = .2), 
  axis.ticks = element_line(linewidth = .2))

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/EnvEffect_anatomicalaxes_barplot.pdf", device = "pdf", dpi = 500, width = 1.55, height = 1.3)

#Calculate empirical enrichment statistic for each S-A axis bin
envSES_effects_SAaxisbins <- spin.data %>% 
                             group_by(SA.axis.bin) %>% 
                             do(mean.tvalue.significant = mean(.$GAM.cov.tvalue, na.rm = TRUE)) %>%
                             unnest(cols = c("mean.tvalue.significant"))

envSES_effects_SAaxisbins

#Calculate null distribution of t-values by bin based on spherical spatial permutations (spins)

## inputs to spatially permute (i.e., spin) with the pre-computed spatial permutation matrix 
x <- spin.data$SA.axis.bin #cortical map 1 to spatially permute: S-A axis bins
y <- spin.data$GAM.cov.tvalue #cortical map 2 to spatially permute: significant envSES-thalamocortical FA t-values
perm.id <- perm.id.full

## spin the empirical data 
nroi = dim(perm.id)[1]  #number of regions
nperm = dim(perm.id)[2] #number of permutations
x.perm = y.perm = array(NA,dim=c(nroi,nperm)) #initialize
  for (r in 1:nperm) {
    for (i in 1:nroi) {
      x.perm[i,r] = x[perm.id[i,r]] #spinning x, spatially permuted bins
      y.perm[i,r] = y[perm.id[i,r]] #spinning y, spatially permuted t-values
    }
  }

## compute mean t-value for each S-A axis bin using spatially permuted x inputs 
### set up spun (x) and empirical (y) df
x.spatialperm <- as.data.frame(x.perm) %>% set_names(sprintf("perm%s", seq(from = 1, to = ncol(x.perm))))
x.spatialperm.empiricaly <- cbind(y, x.spatialperm) %>% as.data.frame()
colnames(x.spatialperm.empiricaly)[1] <- c("empirical.y")

### compute mean t-values for S-A axis bins for all spins
envSES_effects_SAaxisbins_permutedx <- sapply(x.spatialperm.empiricaly[2:ncol(x.spatialperm.empiricaly)], #for all permuted bins
                                              function(p){
                                                x.spatialperm.empiricaly %>% group_by(as.character(p)) %>% #group by permuted S-A axis bins
                                                do(mean.tvalue.significant = mean(.$empirical.y, na.rm = TRUE)) %>% #calculate mean t-value
                                                unnest(cols = c("mean.tvalue.significant")) %>% t()}) %>% 
                                                t() %>% as.data.frame()
                                                 
envSES_effects_SAaxisbins_permutedx <- envSES_effects_SAaxisbins_permutedx %>% mutate_if(is.character, as.numeric) #make it numeric
envSES_effects_SAaxisbins_permutedx <-  envSES_effects_SAaxisbins_permutedx[ , !c(TRUE,FALSE) ] 
colnames(envSES_effects_SAaxisbins_permutedx) <- c("SA.bin1.tvalue","SA.bin2.tvalue", "SA.bin3.tvalue", "SA.bin4.tvalue", "SA.bin5.tvalue") 

## compute mean t-value for each S-A axis bin using spatially permuted y inputs 
### set up spun (y) and empirical (x) df
y.spatialperm <- as.data.frame(y.perm) %>% set_names(sprintf("perm%s", seq(from = 1, to = ncol(y.perm))))
y.spatialperm.empiricalx <- cbind(x, y.spatialperm) %>% as.data.frame()
colnames(y.spatialperm.empiricalx)[1] <- c("empirical.x")

### compute mean t-values for S-A axis bins for all spins
envSES_effects_SAaxisbins_permutedy <- y.spatialperm.empiricalx %>% group_by(empirical.x) %>% #group by empirical S-A axis bins
                                       dplyr::summarize(across(everything(), \(x) mean(x, na.rm = TRUE))) %>% #calculate mean t-value for permuted data
                                       select(-empirical.x) %>% t() %>% as.data.frame() %>%
                                       set_names(c("SA.bin1.tvalue","SA.bin2.tvalue", "SA.bin3.tvalue", "SA.bin4.tvalue", "SA.bin5.tvalue"))

histogram.data <- rbind(envSES_effects_SAaxisbins_permutedx, envSES_effects_SAaxisbins_permutedy)
saveRDS(histogram.data, "/cbica/projects/thalamocortical_development/thalamocortical_results/environment_results/envtvalue.SAbins.spatialnulls.rds")

histogram.data <- readRDS("/cbica/projects/thalamocortical_development/thalamocortical_results/environment_results/envtvalue.SAbins.spatialnulls.rds")
envSES_effects_SAaxisbins_permutedx <- histogram.data[1:10000,]
envSES_effects_SAaxisbins_permutedy <- histogram.data[10001:20000,]
nperm = dim(perm.id.full)[2]

#Function to plot the null distribution of environment effect t-values for a given S-A axis bin and test whether the true t-value is significantly smaller or greater in magnitude than the null (permutation based p)
enveffects_SAbin_enrichment <- function(quintile, enrichment_type){
  
  SAbin.tvalue.empirical <- envSES_effects_SAaxisbins %>% filter(SA.axis.bin == quintile) %>% select(mean.tvalue.significant)
  
  #Histogram
  SA.colorvector <- c("#FFC228", "#FFD16A", "#e9dcf2", "#BE93C4", "#6F1382")
  
  null.tvalue.histogram <- histogram.data %>% select(sprintf("SA.bin%s.tvalue", quintile)) %>%
    ggplot(aes(x = .[,1])) + geom_histogram(fill = c(SA.colorvector)[quintile]) +
    theme_classic() +
    geom_vline(xintercept =  SAbin.tvalue.empirical$mean.tvalue.significant, linewidth = 0.25) +
    xlab("Environment effect (t-value)") +
    ylab("Number of nulls") +
    scale_x_continuous(breaks = c(2, 3, 4, 5), limits = c(1.25, 5.25)) +
    theme(
     legend.position = "none", 
     axis.text.x = element_text(size = 5, family = "Arial", color = c("black")),
     axis.text.y = element_blank(),
     axis.title.x = element_text(size = 5, family ="Arial", color = c("black")),
     axis.title.y = element_text(size = 5, family ="Arial", color = c("black")),
     axis.line.x = element_line(linewidth = .2), 
     axis.line.y = element_blank(),
     axis.ticks = element_line(linewidth = .2))
    
  #Permutation-based p
    if(enrichment_type == "smaller"){
      p.perm.xy <- sum(envSES_effects_SAaxisbins_permutedy[,quintile] < SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
      p.perm.yx <- sum(envSES_effects_SAaxisbins_permutedx[,quintile] < SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
      perm.p <- (p.perm.xy+p.perm.yx)/2
    }
  
  if(enrichment_type == "larger"){
      p.perm.xy <- sum(envSES_effects_SAaxisbins_permutedy[,quintile] > SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
      perm.p <- p.perm.xy
      p.perm.yx <- sum(envSES_effects_SAaxisbins_permutedx[,quintile] > SAbin.tvalue.empirical$mean.tvalue.significant)/nperm
      perm.p <- (p.perm.xy+p.perm.yx)/2
  }
  
  permutation.results <- list(null.tvalue.histogram, perm.p)
  return(permutation.results)
}

enveffects_SAbin_enrichment(quintile = 1, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile1_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(quintile = 2, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile2_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(quintile = 3, enrichment_type = "smaller")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile3_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(quintile = 4, enrichment_type = "larger")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile4_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)

enveffects_SAbin_enrichment(quintile = 5, enrichment_type = "larger")

ggsave(filename = "/cbica/projects/thalamocortical_development/figures/images/Figure7/quintile5_enveffect_spins.pdf", device = "pdf", dpi = 500, width = .95, height =.9)